package com.ltx.aishare.controller;

import com.ltx.aishare.config.AiConfig;
import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.community.model.dashscope.QwenStreamingChatModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import dev.langchain4j.service.TokenStream;
import jakarta.annotation.Resource;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestMethod;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

/**
 * @author 罗添煦
 * @description langChain4j使用ai控制层
 * @create 2025-04-20 下午4:25
 */
@RestController
@RequestMapping("/api/langChain4j")
public class LangChain4jController {

    @Resource
    private QwenChatModel qwenChatModel;
    @Resource
    private QwenStreamingChatModel qwenStreamingChatModel;
    @Resource
    private AiConfig.Assistant assistant;
    @Resource
    private AiConfig.AssistantUnique assistantUnique;
    @Resource
    private AiConfig.AssistantTool assistantTool;

    /**
     * <a href="http://localhost:8081/api/langChain4j/chat?prompt=%E4%BD%A0%E6%98%AF%E8%B0%81">...</a>
     *
     * @param prompt 提示词
     * @return String 返回输出
     */
    @RequestMapping(value = "/chat", method = RequestMethod.GET)
    public String chat(@RequestParam("prompt") String prompt) {
        return qwenChatModel.chat(prompt);
    }

    /**
     * <a href="http://localhost:8081/api/langChain4j/streamingChat?prompt=%E4%BD%A0%E6%98%AF%E8%B0%81">...</a>
     *
     * @param prompt 提示词
     * @return Flux<String>流式返回输出 需要导入pom依赖spring-boot-starter-webflux
     */
    @RequestMapping(value = "/streamingChat", method = RequestMethod.GET, produces = "text/stream;charset=UTF-8")
    public Flux<String> streamingChat(@RequestParam("prompt") String prompt) {
        Flux<String> flux = Flux.create(sink -> {
            qwenStreamingChatModel.chat(prompt, new StreamingChatResponseHandler() {
                @Override
                public void onPartialResponse(String partialResponse) {
                    sink.next(partialResponse);
                }

                @Override
                public void onCompleteResponse(ChatResponse chatResponse) {
                    sink.complete();
                }

                @Override
                public void onError(Throwable throwable) {
                    sink.error(throwable);
                }
            });
        });

        return flux;
    }

    /**
     * <a href="http://localhost:8081/api/langChain4j/memory_chat?prompt=%E6%88%91%E6%98%AF%E7%BD%97%E6%B7%BB%E7%85%A6">...</a>
     *
     * @param prompt 提示词
     * @return String 返回输出
     */
    @RequestMapping(value = "/memory_chat", method = RequestMethod.GET)
    public String memoryChat(@RequestParam(value = "prompt", defaultValue = "我是罗添煦") String prompt) {
        return assistant.chat(prompt);
    }

    /**
     * <a href="http://localhost:8081/api/langChain4j/memory_streaming_chat?prompt=%E6%88%91%E6%98%AF%E8%B0%81">...</a>
     *
     * @param prompt 提示词
     * @return Flux<String>流式返回输出 需要导入pom依赖spring-boot-starter-webflux
     */
    @RequestMapping(value = "/memory_streaming_chat", method = RequestMethod.GET, produces = "text/stream;charset=UTF-8")
    public Flux<String> memoryStreamingChat(@RequestParam(value = "prompt", defaultValue = "我是谁") String prompt) {
        TokenStream stream = assistant.stream(prompt);

        return Flux.create(sink -> {
            stream.onPartialResponse(sink::next)
                    .onCompleteResponse(chatResponse -> {
                        sink.complete();
                    }).onError(sink::error).start();
        });
    }

    /**
     * <a href="http://localhost:8081/api/langChain4j/memory_chat_unique?prompt=%E6%88%91%E6%98%AF%E7%BD%97%E6%B7%BB%E7%85%A6&memoryId=1">...</a>
     * <a href="http://localhost:8081/api/langChain4j/memory_chat_unique?prompt=%E6%88%91%E6%98%AF%E8%B0%81&memoryId=2">...</a>
     *
     * @param prompt   提示词
     * @param memoryId 记忆id
     * @return String 返回输出
     */
    @RequestMapping(value = "/memory_chat_unique", method = RequestMethod.GET)
    public String memoryChatUnique(@RequestParam(value = "prompt", defaultValue = "我是罗添煦") String prompt, @RequestParam(value = "memoryId") int memoryId) {
        return assistantUnique.chat(memoryId, prompt);
    }

    /**
     * <a href="http://localhost:8081/api/langChain4j/memory_streaming_chat_unique?prompt=%E6%88%91%E6%98%AF%E8%B0%81&memoryId=1">...</a>
     *
     * @param prompt   提示词
     * @param memoryId 记忆id
     * @return Flux<String>流式返回输出 需要导入pom依赖spring-boot-starter-webflux
     */
    @RequestMapping(value = "/memory_streaming_chat_unique", method = RequestMethod.GET, produces = "text/stream;charset=UTF-8")
    public Flux<String> memoryStreamingChatUnique(@RequestParam(value = "prompt", defaultValue = "我是谁") String prompt, @RequestParam(value = "memoryId") int memoryId) {
        TokenStream stream = assistantUnique.stream(memoryId, prompt);

        return Flux.create(sink -> {
            stream.onPartialResponse(sink::next)
                    .onCompleteResponse(chatResponse -> {
                        sink.complete();
                    }).onError(sink::error).start();
        });
    }

    /**
     * <a href="http://localhost:8081/api/langChain4j/chat_function?prompt=%E4%B8%AD%E5%92%8C%E5%86%9C%E4%BF%A1%E7%9A%84%E4%B8%BB%E8%A6%81%E4%B8%9A%E5%8A%A1%E6%98%AF%E4%BB%80%E4%B9%88">...</a>
     *
     * @param prompt 提示词
     * @return String 返回输出
     */
    @RequestMapping(value = "/chat_function", method = RequestMethod.GET)
    public String chatFunction(@RequestParam(value = "prompt") String prompt) {
        return assistantTool.chat(prompt);
    }
}
